package cim144.ctdp_array16_shift8_relu
import Chisel._
import chisel3.{Mux, RegInit, RegNext, SInt, UInt, Vec, VecInit, Wire, when}
import chisel3.util.{Cat, RegEnable, log2Ceil}
import scala.io.Source

object mvmcim_shift {
  def apply(custom_mvm:Bool,ismvm:Bool,bitwise_end:Bool,
            adc_range:UInt,out_shift:UInt,isrelu:Bool,push_576:UInt,
            r1:Bits,r2:Bits
           ):(Bool,Bool,UInt)={
    val inst = Module(new mvmcim_shift()).io
    inst.custom_mvm   :=custom_mvm
    inst.ismvm        :=ismvm
    inst.bitwise_end  :=bitwise_end
    inst.adc_range    :=adc_range
    inst.out_shift    :=out_shift
    inst.isrelu       :=isrelu
    inst.push_576     :=push_576
    inst.r1           :=r1
    inst.r2           :=r2
    (inst.cim_done,inst.mvm_done,inst.save_1024)
  }
}


class mvmcim_shift extends Module with SystemConfig{
  val io = new Bundle{
    val custom_mvm  = Input(Bool())
    val ismvm       = Input(Bool())
    val bitwise_end = Input(Bool())
    val adc_range   = Input(UInt(8.W))
    val out_shift   = Input(UInt(8.W))
    val isrelu      = Input(Bool())               // add relu config
    val push_576   = Input(UInt(ROW_NUM.W))
    val r1          = Bits(INPUT,64)
    val r2          = Bits(INPUT,64)

    val cim_done    = Output(Bool())
    val mvm_done    = Output(Bool())
    val save_1024   = Output(UInt((SAVE_MAX_SIZE*CIM_XLEN).W))
  }
  val mvm_start_tmp = RegNext(Mux(io.custom_mvm,true.B,Mux(io.cim_done&io.ismvm,true.B,false.B)))
  val mvm_start     = RegNext(Mux(io.mvm_done,false.B,mvm_start_tmp))
  val rcbd = Wire(new mvm_IO())
  //rs1 |31******************16|15***************0|
  //    |*******row_num********|******col_num*****|
  //rs2 |31*********22|21*********11|10**********0|
  //    |*array index*|**start_row**|**start_col**|
  rcbd.row_begin := RegEnable(io.r2(21,11),0.U,io.custom_mvm)
  rcbd.row_end   := RegEnable(io.r1(31,16)+io.r2(21,11) - 1.U,0.U,io.custom_mvm)
  rcbd.col_begin := RegEnable(io.r2(10, 0),0.U,io.custom_mvm)
  rcbd.col_end   := RegEnable(io.r1(15, 0)+io.r2(10, 0) - 1.U,0.U,io.custom_mvm)
  rcbd.array_index := RegEnable(io.r2(31, 22),0.U,io.custom_mvm)
  // *********** 1. get cim data
  val array_num = 5
  val filelist = "/home/hezf/chipyard_docker/generators/cim144/romfile/filelist.txt"
  val file = Source.fromFile(filelist).getLines().toList
  val array_enable_list = Seq.range(0,array_num).map{i=> rcbd.array_index === i.U}
  val mvm_start_list    = Seq.range(0,array_num).map{i=> mvm_start && array_enable_list(i)}
  val u_cim_out_list    = Seq.range(0,array_num).map{i=> cim_mvm_8bit_shift(mvm_start_list(i),rcbd,io.push_576,io.adc_range,file(i))}
  val cim_done_list     = Seq.range(0,array_num).map{i=> RegNext(u_cim_out_list(i)._1)}
  val cim_result_list   = Seq.range(0,array_num).map{i=>u_cim_out_list(i)._2.asTypeOf(Vec(COL_NUM,SInt(8.W)))}
  io.cim_done  := cim_done_list.reduce( _ || _ )
  //val cim_result = Seq.range(0,array_num).map{i=>Mux(array_enable_list(i),cim_result_list(i),VecInit(Seq.fill(COL_NUM)(0.S(8.W))))}.reduce(_|_)
  val cim_result  = VecInit(cim_result_list)(rcbd.array_index)

//  val array0_enable    = (rcbd.array_index === 0.U)
//  val array1_enable    = (rcbd.array_index === 1.U)
//  val mvm_start_array0 = mvm_start && array0_enable
//  val mvm_start_array1 = mvm_start && array1_enable
//  val u0_cim_out     = cim_mvm_8bit_shift(mvm_start_array0,rcbd,io.push_576,io.adc_range,
//                      "/home/hezf/chipyard_docker/generators/cim144/mlp_weight_rom.txt")
//  val cim_done0  = RegNext(u0_cim_out._1)
//  val u1_cim_out     = cim_mvm_8bit_shift(mvm_start_array1,rcbd,io.push_576,io.adc_range,"None")
//  val cim_done1  = RegNext(u1_cim_out._1)
//  io.cim_done    := cim_done0 | cim_done1
//  val cim_result0 = u0_cim_out._2.asTypeOf(Vec(COL_NUM,SInt(8.W)))
//  val cim_result1 = u1_cim_out._2.asTypeOf(Vec(COL_NUM,SInt(8.W)))
//  val cim_result  = Mux(array0_enable,cim_result0,
//                    Mux(array1_enable,cim_result1, VecInit(Seq.fill(COL_NUM)(0.S(8.W)))))

  // ************ 2. Acc
  val mvm_tmp         = RegInit(VecInit(Seq.fill(COL_NUM)(0.S(16.W))))
  val mvm_tmp_shift   = VecInit(Seq.fill(COL_NUM)(0.S(16.W)))
  val mvm_result      = VecInit(Seq.fill(COL_NUM)(0.S(8.W)))
  val mvm_result_unrelu = VecInit(Seq.fill(COL_NUM)(0.S(8.W)))
  // +++++ADD shift_num
  val shift_num       = RegInit(0.U(4.W))
  for(i<-0 until COL_NUM) {
    when(io.custom_mvm) {
      mvm_tmp(i) := 0.S
      shift_num  := 0.U  //++++++++
    }.elsewhen(io.cim_done) {
      mvm_tmp(i) := mvm_tmp(i) + (cim_result(i) << shift_num)  //+++++++++
      shift_num := shift_num + 1.U  //++++++++++
    }
    mvm_tmp_shift(i) := mvm_tmp(i) >> io.out_shift
    mvm_result_unrelu(i) := Mux(mvm_tmp_shift(i) >=  127.S,  127.S,
                            Mux(mvm_tmp_shift(i) <= -127.S, -127.S,
                                mvm_tmp_shift(i)))
    mvm_result(i) := Mux(io.isrelu,Mux(mvm_result_unrelu(i)(7),0.S,mvm_result_unrelu(i))
                                  ,mvm_result_unrelu(i))
  }
  io.mvm_done := io.ismvm & io.bitwise_end
  //io.save_1024 = mvm_result.asTypeOf(Vec(SAVE_MAX_SIZE,UInt(CIM_XLEN.W)))
  io.save_1024 := mvm_result.asUInt()
}
